package nmap

import (
	"bytes"
	"context"
	"crypto/tls"
	"fmt"
	hubur "gitee.com/wudicaidou/menciis-hubur"
	logx "gitee.com/wudicaidou/menciis-logx"
	"github.com/dlclark/regexp2"
	"github.com/pkg/errors"
	"golang.org/x/net/proxy"
	"net"
	"net/url"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync/atomic"
	"time"
	"unicode"
)

type versionTask struct {
	host    string
	port    int
	addr    string
	network string
	proto   TransportProto
	tls     bool
	isOpen  bool
}

type VersionDetectorConfig struct {
	DialTimeout           time.Duration
	ReadTimeout           time.Duration
	Rarity                int
	Trace                 bool // debug 使用, 会打印每一步的详情
	ExcludeDangerousPorts bool // 是否排除部分危险端口，比如打印端口，往这些端口发数据会被打印出来
	Proxy                 string
}

type VersionDetector struct {
	*ProbeParseResult
	SrvMap   map[string][]*Probe // Match.ServiceName -> []{*Probe}, 处理 softmatch 的情况
	ProbeMap map[string]*Probe   // Probe.Name -> Probe, 处理 fallback

	config      *VersionDetectorConfig
	logger      *logx.Logger
	dialer      *net.Dialer
	proxyDialer proxy.Dialer

	connectionCount uint64
	defaultService  map[int]string
}

func NewVersionDetector(config *VersionDetectorConfig, probes *ProbeParseResult) *VersionDetector {
	if len(probes.Probes) == 0 {
		panic("parse result can't be nil")
	}
	srvMap := make(map[string][]*Probe)
	probeMap := make(map[string]*Probe)
	for _, p := range probes.Probes {
		probeMap[p.Name] = p
		for _, m := range p.Matches {
			srvMap[m.ServiceName] = append(srvMap[m.ServiceName], p)
		}
	}

	// 去重
	for k, ps := range srvMap {
		sort.Slice(ps, func(i, j int) bool {
			return ps[i].Index < ps[j].Index
		})
		var newPs []*Probe
		lastIndex := -1
		for _, p := range ps {
			if p.Index != lastIndex {
				newPs = append(newPs, p)
				lastIndex = p.Index
			}
		}
		srvMap[k] = newPs
	}

	return &VersionDetector{
		defaultService:   GetDefaultServices(),
		ProbeParseResult: probes,
		SrvMap:           srvMap,
		ProbeMap:         probeMap,
		config:           config,
		dialer: &net.Dialer{
			Timeout: config.DialTimeout,
		},
	}
}

func (v *VersionDetector) SetLogger(l *logx.Logger) {
	v.logger = l
}

func (v *VersionDetector) SetProxy(proxyUrl string) {
	u, err := url.Parse(proxyUrl)
	if err == nil {
		dialer, err := proxy.FromURL(u, proxy.Direct)
		if err == nil {
			v.proxyDialer = dialer
		}
	}
}

type tracedConn struct {
	net.Conn
}

func (t *tracedConn) Read(b []byte) (n int, err error) {
	n, err = t.Conn.Read(b)
	//fmt.Printf("\nread data from %s\n", t.RemoteAddr())
	//fmt.Print(hex.Dump(b[:n]))
	return n, err
}

func (t *tracedConn) Write(b []byte) (n int, err error) {
	//fmt.Printf("\nwrite data to %s\n", t.RemoteAddr())
	//fmt.Print(hex.Dump(b))
	return t.Conn.Write(b)
}

func (v *VersionDetector) getConn(ctx context.Context, network string, addr string, ssl bool) (net.Conn, error) {
	atomic.AddUint64(&v.connectionCount, 1)
	newCtx, cancel := context.WithTimeout(ctx, v.config.DialTimeout)
	defer cancel()

	var conn net.Conn
	var err error
	if v.proxyDialer != nil {
		proxyDialer, ok := v.proxyDialer.(proxy.ContextDialer)
		if !ok {
			return nil, errors.New("invalid proxy dialer")
		}
		conn, err = proxyDialer.DialContext(ctx, network, addr)
		if err != nil {
			return nil, err
		}
	} else {
		conn, err = v.dialer.DialContext(newCtx, network, addr)
		if err != nil {
			return nil, err
		}
	}

	if ssl {
		conn = tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
		_ = conn.SetDeadline(time.Now().Add(v.config.ReadTimeout))
		if err := conn.(*tls.Conn).Handshake(); err != nil {
			_ = conn.Close()
			return nil, err
		}
	}
	if v.config.Trace {
		conn = &tracedConn{conn}
	}
	return conn, nil
}

func (v *VersionDetector) GetConnectCount() uint64 {
	return atomic.LoadUint64(&v.connectionCount)
}

func (v *VersionDetector) Detect(ctx context.Context, host string, port int, proto TransportProto) (service *ServiceInfo, isOpen bool, err error) {
	service, isOpen, err = v.DetectReal(ctx, host, port, proto)
	if err == nil && isOpen && service == nil {
		if v, ok := v.defaultService[port]; ok {
			service = &ServiceInfo{
				ServiceVersion: ServiceVersion{
					ServiceName: v,
				},
				Host:  host,
				Port:  port,
				Proto: proto,
			}
			err = nil
		}
	}

	return service, isOpen, err
}

func (v *VersionDetector) DetectReal(ctx context.Context, host string, port int, proto TransportProto) (service *ServiceInfo, isOpen bool, err error) {
	if v.config.ExcludeDangerousPorts {
		if proto == TCP {
			if v.ExcludePorts.InTCP(port) {
				return nil, false, fmt.Errorf("%d is in dangrous port list and will be ignored", port)
			}
		} else {
			if v.ExcludePorts.InUDP(port) {
				return nil, false, fmt.Errorf("%d is in dangrous port list and will be ignored", port)
			}
		}
	}
	task := &versionTask{
		host:  host,
		port:  port,
		proto: proto,
		addr:  net.JoinHostPort(host, strconv.Itoa(port)),
		tls:   false,
	}
	if proto == TCP {
		task.network = "tcp"
	} else {
		task.network = "udp"
	}
	srv, err := v.detect(ctx, task)
	if err != nil {
		return srv, task.isOpen, err
	}
	if srv == nil {
		if task.proto == TCP {
			return nil, true, nil
		} else {
			return nil, false, nil
		}
	}
	if srv.ServiceName != "ssl" {
		return srv, true, nil
	}
	srv.SSL = true
	//v.logger.Infof("service %s is ssl, retry with ssl", task.addr)
	task.tls = true
	sslSrv, err := v.detect(ctx, task)
	if err != nil {
		return nil, true, err
	}
	if sslSrv != nil {
		return sslSrv, true, nil
	}
	//v.logger.Infof("service %s is ssl, but can not match other service", task.addr)
	return srv, true, nil
}

func (v *VersionDetector) detect(ctx context.Context, task *versionTask) (*ServiceInfo, error) {
	if task.proto == TCP {
		// detect null Probe
		version, err := v.doProbe(ctx, v.Probes[0], task, false)
		if err != nil {
			return nil, err
		}
		// 如果 null 没有报错，认为开放
		task.isOpen = true
		if version != nil {
			return version, nil
		}
	}

	// ports 命中，无视 rarity
	//v.logger.Infof("check probes that matches the host(%s), ssl(%v), port(%d) ", task.host, task.tls, task.port)
	for _, probe := range v.Probes[1:] {
		if probe.Proto != task.proto {
			continue
		}
		if v.portInProbe(task, probe) {
			version, err := v.doProbe(ctx, probe, task, false)
			if err != nil {
				return nil, fmt.Errorf("probe %v return error, %s", probe, err)
			}
			if version != nil {
				return version, nil
			}
		}
	}

	// all other ports
	v.logger.Verbosef("check all the rest probes that do not Match the ports for %s", task.addr)

	for _, probe := range v.Probes[1:] {
		if probe.Proto != task.proto || probe.Rarity > v.config.Rarity {
			continue
		}
		if v.portInProbe(task, probe) {
			continue
		}
		version, err := v.doProbe(ctx, probe, task, false)
		if err != nil {
			return nil, fmt.Errorf("probe %v return error, %s", probe, err)
		}
		if version != nil {
			return version, nil
		}
	}
	return nil, nil
}

func (v *VersionDetector) portInProbe(task *versionTask, probe *Probe) bool {
	if task.proto == TCP {
		b := (task.tls && probe.SSLPorts != nil && probe.SSLPorts.InTCP(task.port) && probe.Rarity != 9) ||
			(!task.tls && probe.Ports != nil && probe.Ports.InTCP(task.port) && probe.Rarity != 9)
		//fmt.Println(b, probe.Rarity)
		return b
	} else {
		return (task.tls && probe.SSLPorts != nil && probe.SSLPorts.InUDP(task.port) && probe.Rarity != 9) ||
			(!task.tls && probe.Ports != nil && probe.Ports.InUDP(task.port) && probe.Rarity != 9)
	}
}

func (v *VersionDetector) serviceInfoFromVersion(version *ServiceVersion, task *versionTask, banner []byte) *ServiceInfo {
	return &ServiceInfo{
		Host:           task.host,
		Port:           task.port,
		Proto:          task.proto,
		SSL:            task.tls,
		Banner:         hubur.EscapeInvalidUTF8Byte(banner),
		ServiceVersion: *version,
	}
}

func (v *VersionDetector) doProbe(ctx context.Context, p *Probe, task *versionTask, allowSoft bool) (*ServiceInfo, error) {
	v.logger.Verbosef("running %s Probe for %s", p, task.addr)

	conn, err := v.getConn(ctx, task.network, task.addr, task.tls)
	if err != nil {
		return nil, err
	}
	defer func() {
		_ = conn.Close()
	}()
	_ = conn.SetWriteDeadline(time.Now().Add(v.config.DialTimeout))
	_, err = conn.Write(p.Payload)
	if err != nil {
		return nil, err
	}
	recv := v.readUntilTimeout(conn, v.getReadTimeout(p))
	if len(recv) != 0 {
		version, soft := v.doMatch(p, recv, task)
		if version != nil {
			if !soft {
				return v.serviceInfoFromVersion(version, task, recv), nil
			}
			// 及早关闭, 忽略 err
			_ = conn.Close()
			newVersion, err := v.softMatch(ctx, p, version.ServiceName, task)
			if err != nil {
				return nil, err
			}
			if newVersion == nil {
				// 如果 softmatch 失败了，结束逻辑，使用 softmatch 的结果
				return v.serviceInfoFromVersion(version, task, recv), nil
			}
			return newVersion, nil
		}
		fallbacks := p.Fallback
		if task.proto == TCP {
			fallbacks = append(fallbacks, ProbeNULL)
		}
		if len(fallbacks) == 0 {
			return nil, nil
		}
		// v.logger.Debugf("running fallback %v for %s in %s", fallbacks, task.addr, p.Name)
		for _, fallback := range fallbacks {
			p := v.ProbeMap[fallback]
			if p == nil {
				v.logger.Errorf("can't found the fallback Probe %s", fallback)
				continue
			}
			version, soft := v.doMatch(p, recv, task)
			if version != nil {
				if !soft || allowSoft {
					return v.serviceInfoFromVersion(version, task, recv), nil
				}
				// 及早关闭, 忽略 err
				_ = conn.Close()
				newVersion, err := v.softMatch(ctx, p, version.ServiceName, task)
				if err != nil {
					return nil, err
				}
				if newVersion == nil {
					// 如果 softmatch 失败了，结束逻辑，使用 softmatch 的结果
					return v.serviceInfoFromVersion(version, task, recv), nil
				}
				return newVersion, nil
			}
		}
	}
	return nil, nil
}

func (v *VersionDetector) softMatch(ctx context.Context, sourceProbe *Probe, srvName string, task *versionTask) (*ServiceInfo, error) {
	//v.logger.Infof("running softmatch %s for %s", srvName, task.addr)

	for _, p := range v.SrvMap[srvName] {
		// 有序的，前面的都跑过了
		if p.Index <= sourceProbe.Index {
			continue
		}
		// 从 softmatch 出来的缺符合 port 规则，说明这条 probe 没有在之前的阶段跑出来，这一般是有 bug
		if v.portInProbe(task, p) && v.config.Trace {
			v.logger.Warnf("port matches but no rule matches, may be a bug: %s %s:%d", task.network, task.host, task.port)
		}

		version, err := v.doProbe(ctx, p, task, true)
		if err != nil {
			return nil, err
		}
		if version != nil {
			return version, nil
		}
	}
	return nil, nil
}

func (v *VersionDetector) doMatch(p *Probe, recv []byte, task *versionTask) (*ServiceVersion, bool) {
	recvRune := make([]rune, len(recv))
	for i, c := range recv {
		recvRune[i] = rune(c)
	}

	extractVersion := func(m *Match, result *regexp2.Match) *ServiceVersion {
		// 绝不能对 Match 本身进行修改, 否则会有竞争
		ret := &ServiceVersion{
			//ServiceName: m.ServiceName,
			ServiceName: FixProtocol(m.ServiceName),
			Product:     m.Product,
			Version:     m.Version,
			Info:        m.Info,
			Hostname:    m.Hostname,
			DeviceType:  m.DeviceType,
			OS:          m.OS,
			CPE:         append(m.CPE[:0:0], m.CPE...),
		}
		ret.Product = v.processSubMatch(ret.Product, result)
		ret.Version = v.processSubMatch(ret.Version, result)
		ret.Info = v.processSubMatch(ret.Info, result)
		ret.Hostname = v.processSubMatch(ret.Hostname, result)
		ret.DeviceType = v.processSubMatch(ret.DeviceType, result)
		ret.OS = v.processSubMatch(ret.OS, result)
		for i, cpe := range ret.CPE {
			ret.CPE[i] = v.processSubMatch(cpe, result)
		}
		return ret
	}
	for _, match := range p.Matches {
		m := match
		result, err := match.Rule.FindRunesMatch(recvRune)
		if err != nil {
			if !strings.Contains(err.Error(), "timeout") {
				v.logger.Error(err)
			}
			continue
		}
		if result == nil || result.GroupCount() == 0 {
			continue
		}
		//v.logger.Infof("service %s match probe: %s with %s, soft: %v,rule: %v, rarity: %v", task.addr, p.Name, m.ServiceName, m.Soft, match.Rule.String(), p.Rarity)
		return extractVersion(m, result), m.Soft
	}

	return nil, false
}

var regexpFirstNum = regexp.MustCompile(`^\d`)

func FixProtocol(oldProtocol string) string {
	//进行最后输出修饰
	if oldProtocol == "ssl/http" {
		return "https"
	}
	if oldProtocol == "http-proxy" {
		return "http"
	}
	if oldProtocol == "ms-wbt-server" {
		return "rdp"
	}
	if oldProtocol == "microsoft-ds" {
		return "smb"
	}
	if oldProtocol == "netbios-ssn" {
		return "netbios"
	}
	if oldProtocol == "oracle-tns" {
		return "oracle"
	}
	if oldProtocol == "msrpc" {
		return "rpc"
	}
	if oldProtocol == "ms-sql-s" {
		return "mssql"
	}
	if oldProtocol == "domain" {
		return "dns"
	}
	if oldProtocol == "svnserve" {
		return "svn"
	}
	if oldProtocol == "ibm-db2" {
		return "db2"
	}
	if oldProtocol == "socks-proxy" {
		return "socks5"
	}
	if len(oldProtocol) > 4 {
		if oldProtocol[:4] == "ssl/" {
			return oldProtocol[4:] + "-ssl"
		}
	}
	if regexpFirstNum.MatchString(oldProtocol) {
		oldProtocol = "S" + oldProtocol
	}
	oldProtocol = strings.ReplaceAll(oldProtocol, "_", "-")
	return oldProtocol
}

const (
	exprStateData = iota
	exprStateNameStart
	exprStateInArgumentList
	exprStateInArgument
	exprStateReplaceIndex
)

// 处理四种情况:
// $0
// $P(i)
// $SUBST(i,"","")
// $I(i, ">/<")
func (v *VersionDetector) processSubMatch(raw string, result *regexp2.Match) string {
	state := exprStateData
	var fnName string
	var buf bytes.Buffer
	var index int
	var args []interface{}
	var lastArg interface{}
	var stringArg string
	data := []byte(raw)
OuterLoop:
	for _, r := range data {
		switch state {
		case exprStateData:
			if r == '$' {
				state = exprStateNameStart
				continue
			}
			buf.WriteByte(r)
		case exprStateNameStart:
			if r > '0' && r < '9' {
				index = index*10 + int(r-'0')
				state = exprStateReplaceIndex
				continue
			}
			if r == '(' {
				state = exprStateInArgumentList
				continue
			}
			fnName += string(r)
		case exprStateReplaceIndex:
			if r > '0' && r < '9' {
				index = index*10 + int(r-'0')
				continue
			}
			group := result.GroupByNumber(index)
			if group != nil {
				buf.WriteString(group.String())
			} else {
				buf.WriteString(fmt.Sprintf("$%d", index))
				v.logger.Errorf("group index error, matched data: %s, index: %d", result.String(), index)
			}
			index = 0
			if r == '$' {
				state = exprStateNameStart
				continue
			}
			buf.WriteByte(r)
			state = exprStateData
		case exprStateInArgumentList:
			if r == ' ' {
				continue
			}
			if r > '0' && r < '9' {
				index = index*10 + int(r-'0')
				lastArg = index
				continue
			}
			switch r {
			case '"':
				state = exprStateInArgument
			case ',':
				index = 0
				args = append(args, lastArg)
			case ')':
				args = append(args, lastArg)
				var err error
				switch fnName {
				case "P":
					err = v.funcToPrintable(args, result, &buf)
				case "SUBST":
					err = v.funcSubString(args, result, &buf)
				case "I":
					err = v.funcUnpackToInt(args, result, &buf)
				default:
					v.logger.Errorf("unknow fnName: %s", fnName)
					buf.WriteString(fmt.Sprintf("$%s(%#v)", fnName, args))
				}
				if err != nil {
					v.logger.Errorf("can't process the function: %s, args: %+v", fnName, args)
					break OuterLoop
				}
				args = args[:0]
				fnName = ""
				index = 0
				stringArg = ""
				state = exprStateData
			}
		case exprStateInArgument:
			if r == '"' {
				lastArg = stringArg
				stringArg = ""
				state = exprStateInArgumentList
				continue
			}
			stringArg += string(r)
		}
	}
	// $1 这种没有结束标记, 如果正好结束，会导致最后一个 $i 没有被替换，在最后处理下
	if index != 0 {
		group := result.GroupByNumber(index)
		if group != nil {
			buf.WriteString(group.String())
		} else {
			buf.WriteString(fmt.Sprintf("$%d", index))
			v.logger.Errorf("group date index error, matched data: %s, index: %d", result.String(), index)
		}
	}
	return buf.String()
}

// 实现 nmap 要求的 $P(i) 函数，去除不可打印的字符
func (v *VersionDetector) funcToPrintable(args []interface{}, result *regexp2.Match, buf *bytes.Buffer) error {
	if len(args) != 1 {
		return fmt.Errorf("function $P args error, %v", args)
	}
	i, ok := args[0].(int)
	if !ok {
		return fmt.Errorf("function $P args error, %v", args)
	}
	m := result.GroupByNumber(i)
	if m == nil {
		return fmt.Errorf("function $P matches index err, %d, count: %d", i, result.GroupCount())
	}
	for _, r := range m.String() {
		if unicode.IsPrint(r) {
			buf.WriteRune(r)
		}
	}
	return nil
}

// 实现 nmap 要求的 $SUBST(i,"","") 函数
func (v *VersionDetector) funcSubString(args []interface{}, result *regexp2.Match, buf *bytes.Buffer) error {
	if len(args) != 3 {
		return fmt.Errorf("function $SUBST args error, %v", args)
	}
	i, ok1 := args[0].(int)
	oldS, ok2 := args[1].(string)
	newS, ok3 := args[2].(string)
	if !ok1 || !ok2 || !ok3 {
		return fmt.Errorf("function $SUBST args error, %v", args)
	}
	m := result.GroupByNumber(i)
	if m == nil {
		return fmt.Errorf("function $P matches index err, %d, count: %d", i, result.GroupCount())
	}
	buf.WriteString(strings.ReplaceAll(m.String(), oldS, newS))
	return nil
}

// 实现 nmap 要求的 $I(i) 函数
func (v *VersionDetector) funcUnpackToInt(args []interface{}, result *regexp2.Match, buf *bytes.Buffer) error {
	if len(args) != 2 {
		return fmt.Errorf("function $I args error, %v", args)
	}
	i, ok1 := args[0].(int)
	edian, ok2 := args[1].(string)
	if !ok1 || !ok2 {
		return fmt.Errorf("function $I args error, %v", args)
	}
	m := result.GroupByNumber(i)
	if m == nil {
		return fmt.Errorf("function $P matches index err, %d, count: %d", i, result.GroupCount())
	}

	data := []byte(m.String())
	if len(data) == 0 {
		return fmt.Errorf("empty bytes to unpack")
	}

	if len(data) > 8 {
		v.logger.Warnf("binary to unpack is too long, %#x", data)
		data = data[:8]
	}
	var ret uint64
	var shift uint64
	// 不要直接用 binary.*Endian, 可能存在非 2^n 的情况, 会 panic
	if edian == "<" {
		for _, b := range data {
			ret |= uint64(b) << shift
			shift += 8
		}
	} else if edian == ">" {
		for i := len(data) - 1; i >= 0; i-- {
			ret |= uint64(data[i]) << shift
			shift += 8
		}
	} else {
		return fmt.Errorf("unknow byte order: %s", edian)
	}

	buf.WriteString(strconv.FormatUint(ret, 10))
	return nil
}

func (v *VersionDetector) getReadTimeout(p *Probe) time.Duration {
	if p.WaitTimeout != 0 {
		return p.WaitTimeout
	}
	return v.config.ReadTimeout
}

func (v *VersionDetector) readUntilTimeout(conn net.Conn, timeout time.Duration) []byte {
	var ret []byte
	for {
		m := make([]byte, 2048)
		_ = conn.SetDeadline(time.Now().Add(timeout))
		n, err := conn.Read(m)
		if n != 2048 || err != nil {
			if len(ret) == 0 {
				return m[:n]
			}
			ret = append(ret, m[:n]...)
			return ret
		}
		ret = append(ret, m[:n]...)
	}
}
